package com.xiashitech.agent.instrumentation.spring.security.inst;

import com.xiashitech.agent.cache.status.CacheStatus;
import com.xiashitech.agent.dto.UserRequest;
import com.xiashitech.agent.instrumentation.helper.calldepth.XSSecurityContextRepositoryCallDepth;
import com.xiashitech.agent.instrumentation.spring.security.context.SecurityRequestContext;
import com.xiashitech.agent.utils.DurationPrinter;
import io.opentelemetry.javaagent.bootstrap.CallDepth;
import io.opentelemetry.javaagent.extension.instrumentation.TypeInstrumentation;
import io.opentelemetry.javaagent.extension.instrumentation.TypeTransformer;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.description.type.TypeDescription;
import net.bytebuddy.matcher.ElementMatcher;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.context.SecurityContext;
import org.springframework.security.web.authentication.WebAuthenticationDetails;

import static io.opentelemetry.javaagent.extension.matcher.AgentElementMatchers.hasSuperType;
import static net.bytebuddy.matcher.ElementMatchers.*;

public class XSSecurityContextRepositoryInstrumentation implements TypeInstrumentation {
    @Override
    public ElementMatcher<TypeDescription> typeMatcher() {
        return hasSuperType(named("org.springframework.security.web.context.SecurityContextRepository"));
    }

    @Override
    public void transform(TypeTransformer transformer) {
        transformer.applyAdviceToMethod(
            isMethod()
                .and(isPublic())
                .and(named("loadContext"))
                .and(takesArgument(0, named("org.springframework.security.web.context.HttpRequestResponseHolder")))
                .and(takesArguments(1))
                .and(returns(named("org.springframework.security.core.context.SecurityContext"))),
                XSSecurityContextRepositoryInstrumentation.class.getName() + "$ContextAdvice");
    }
    public static class ContextAdvice {

        @Advice.OnMethodEnter(suppress = Throwable.class)
        public static void onEnter(
                @Advice.Local("otelCallDepth") CallDepth callDepth) {
            DurationPrinter durationPrinter = new DurationPrinter("XSSecurityContextRepositoryInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            callDepth = CallDepth.forClass(XSSecurityContextRepositoryCallDepth.class);
            if (callDepth.getAndIncrement() > 0) {
                return;
            }
            durationPrinter.printElapseDuration("ContextAdvice b");
        }

        @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
        public static void onExit(
                @Advice.Return SecurityContext securityContext,
                @Advice.Local("otelCallDepth") CallDepth callDepth) {
            DurationPrinter durationPrinter = new DurationPrinter("XSSecurityContextRepositoryInstrumentation");
            if(CacheStatus.cannotRun()) {
                return;
            }
            if (callDepth.decrementAndGet() > 0) {
                return;
            }
            UserRequest securityUserRequest = SecurityRequestContext.getSecurityUserRequest();
            if(securityUserRequest != null) {
                Authentication authentication = securityContext.getAuthentication();
//                    Optional<UserRequest> userRequestOptional = Java8Bridge.executeRuleForUserInfoMatchRule(authentication);
//                    if(userRequestOptional.isPresent()) {
//                        securityUserRequest.setUserid(userRequestOptional.get().getUserid());
//                    }
                if(authentication.getDetails() instanceof WebAuthenticationDetails) {
                    WebAuthenticationDetails webAuthenticationDetails = (WebAuthenticationDetails)authentication.getDetails();
                    securityUserRequest.setSessionId(webAuthenticationDetails.getSessionId());
                }
                securityUserRequest.setLoginName(authentication.getName());
            }
            durationPrinter.printElapseDuration("ContextAdvice e");
        }
    }
}
